import numpy as np
from basic_models.models.TranDGN import TranDGN, TranDGNWithoutMulti, TranDGNWithoutProcessor
from basic_models.models.DGNS import DGN_R2
from basic_models.models.Gacpn import GACPN_R2, GACPN_R3


def params_count(model):
    """mode
    Compute the number of   parameters.
    Args:
        model (model): model to count   the number of parameters.
    """

    return np.sum([p.numel() for p in model.parameters()]).item()


t = TranDGN(10, 20, 128, 5, 9)
d = DGN_R2(10, 20, 256, 5)
g = GACPN_R3(10, 20, 256, 5)
tran_without_multi = TranDGNWithoutMulti(10, 20, 128, 5, 9)
tran_without_processor = TranDGNWithoutProcessor(10, 20, 128, 5, 9)
print(f'TranDGN:{params_count(t)}\n'
      f'DGN:{params_count(d)}\n'
      f'MAGIC:{params_count(g)}\n'
      f'TranWithoutMulti:{params_count(tran_without_multi)}\n'
      f'tran_without_processor:{params_count(tran_without_processor)}')
